/**@addtogroup Filter Steam limiter filter hook DLL.
 * @{@file
 *
 * This defines a data structure to be used for managing replacement HTTP
 * content to be used when certain URL patterns are seen.
 *
 * @author Nigel Bree <nigel.bree@gmail.com>
 *
 * Copyright (C) 2013 Nigel Bree; All Rights Reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 
 * Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 * 
 * Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*
 * The fundamental idea here is that we maintain a small set of data for
 * replacement documents (or document templates; although there doesn't seem
 * to be an immediate need for it here, obviously it would be desirable in some
 * contexts for the replacement document to be generated by template expansion)
 * along with a simple kind of state-tracking structure for socket handles.
 *
 * Since replacement events are expected to be rare, there is generally only
 * likely to be a single outstanding one, and generally it will be consumed in
 * a single read call, so a simple linked list should do for matching things.
 *
 * When a replacement URL is detected, the replacement data structure is set up
 * and the caller's request is discarded so that the connected peer does not
 * see the request being replaced at all (which saves the need to edit the
 * replaced response out of the data stream). In general, requests are sent as
 * a single write call - the problems caused by the behaviour of socket reads
 * (even in synchronous mode, even a single byte of data available causes an
 * immediate return to the caller) tend not to affect write processing.
 *
 * Then, on read processing if a replacement is set to occur for a socket
 * handle the replacement data is written to the user buffer (typically all of
 * it, although for robustness it's best to allow it to be drained in parts),
 * and once the replacement is drained then we can return to the caller.
 *
 * The final detail is the question of where the replacement documents are read
 * from, and when they are read. Potentially the reading could take place at
 * the time a filter rule is established, which would be most prudent if the
 * replacement is taken from a disk file and is thus likely to be expensive and
 * high-latency. However it seems more product to read the document from the
 * system registry if the document is likely to be small (as it is for Steam
 * content filtering) which means that the data encoding is better-specified
 * and due to aggressive system caching of the registry hive data is very
 * likely to be low-latency to access.
 */

#define WIN32_LEAN_AND_MEAN     1
#include <windows.h>
#include <winsock2.h>

#include "replace.h"

/**
 * Cliche for measuring array lengths, to avoid mistakes with sizeof ().
 */

#define ARRAY_LENGTH(x) (sizeof (x) / sizeof (* (x)))

/**
 * Base object for implementing socket state tracking.
 *
 * This is set up for binding sockets and event handles so that applications
 * get notified when a socket becomes readable, since the lack of true AIO in
 * classic sockets requires separate eventing mechanisms (from the simple to
 * the absurd as in epoll (), which is essentially socket-specific as well as
 * baroque and hard to use compared to a universal AIO model).
 */

class SocketTrack {
        /*
         * Friend-ing an entire template is not something you see often, but
         * it's been a part of the language for a while.
         */

        template <class T>
        friend class SocketList;

private:
        SocketTrack   * m_next;
        SocketTrack   * m_prev;
        SOCKET          m_handle;

public:
        WSAEVENT        m_event;

private:
        /* NOCOPY */    SocketTrack (const SocketTrack &);
        void            operator = (const SocketTrack &);

public:
                        SocketTrack (SOCKET handle);

        SOCKET          handle (void) const { return m_handle; }

static  void          * operator new (size_t length) throw ();
static  void          * operator new (size_t length, void * mem) throw ();
static  void            operator delete (void * mem) throw ();
};

/**
 * Regular replacement new, non-throwing.
 */

/* static */
void * SocketTrack :: operator new (size_t length) throw () {
        return HeapAlloc (GetProcessHeap (), 0, length);
}

/**
 * Non-throwing placement new, generally used with manual calls to the plain
 * operator function to allocate variable-sized memory blocks.
 */

/* static */
void * SocketTrack :: operator new (size_t length, void * mem) throw () {
        return mem;
}

/**
 * Trivial deallocator to match the replacement new.
 */

/* static */
void SocketTrack :: operator delete (void * mem) throw () {
        HeapFree (GetProcessHeap (), 0, mem);
}

/**
 * Trivial constructor.
 */

SocketTrack :: SocketTrack (SOCKET handle) : m_next (0), m_prev (0),
                m_handle (handle), m_event (0) {
}

/**
 * Simple list-holder.
 *
 * Intrusive lists have been a part of C++ forever, since they were the common
 * thing done in C for many years before. However, in early C++ what was done
 * in most programs was to convert a common node structure and extend it using
 * virtual functions.
 *
 * Although this could be done to be typesafe, it had problems not the least of
 * which was typesafety required constantly injecting new virtual functions
 * into the base classes or the moral equivalent (e.g. COM-style QueryInterface
 * which was elegant but not being integrated into the language, lots of work).
 *
 * Templates in the 1990 ARM promised to help with this by introducing some
 * parametric genericity, but being mostly a system pf complex hygenic macro-
 * expanders the quality of implementation was low and remained low for well
 * into the 2000s and even then tended to cause code explosion. Finally, around
 * the late 2000's quality implementations of templates style could really
 * rival the old pre-1990s one in total efficiency thanks to function-level
 * linking and deduplication (although that's not entirely free either, as it
 * can make for a disorienting debug experience).
 */

template <class T>
struct SocketList {
typedef CRITICAL_SECTION      Mutex;

        Mutex           m_lock [1];
        T             * m_head;

                        SocketList ();
                      ~ SocketList ();

        void            free (void);

        void            add (T & item);
        T             * find (SOCKET handle);
        void            remove (T * item, bool free = false);
        void            remove (SOCKET handle);
};

/**
 * Initialize an empty list.
 */

template <class T>
SocketList<T> :: SocketList () : m_head (0) {
        InitializeCriticalSection (m_lock);
}

/**
 * Deinitialize the list and locking structure.
 */

template <class T>
SocketList<T> :: ~ SocketList () {
        free ();

        DeleteCriticalSection (m_lock);
}

/**
 * Free all the list members, deallocating them.
 */

template <class T>
void SocketList<T> :: free (void) {
        EnterCriticalSection (m_lock);

        while (m_head != 0)
                remove (m_head, true);

        LeaveCriticalSection (m_lock);
}

/**
 * Add a new item.
 */

template <class T>
void SocketList<T> :: add (T & item) {
        EnterCriticalSection (m_lock);

        item.m_next = m_head;
        item.m_prev = 0;
        m_head = & item;

        LeaveCriticalSection (m_lock);
}

/**
 * Find an item.
 *
 * In a more general template, this would be a member template so that the
 * key type and comparison function could be parameterized, but this is just
 * a simplified example of the general style so I haven't bothered.
 *
 * The key reason this is here is to show a clean example of something memory-
 * light, which nothing in the modern STL can really claim to be (even for the
 * few good parts of classic STL, the retrofitting of exceptions into the STL
 * containers meant it's usually easier to ditch them in memory-constrained or
 * code-size-limited environments).
 */

template <class T>
T * SocketList<T> :: find (SOCKET handle) {
        EnterCriticalSection (m_lock);

        T             * scan = m_head;
        while (scan != 0) {
                if (scan->m_handle == handle)
                        break;

                scan = (T *) scan->m_next;
        }

        LeaveCriticalSection (m_lock);
        return scan;
}

/**
 * Remove an item from the list, optionally freeing it.
 */

template <class T>
void SocketList<T> :: remove (T * item, bool free) {
        EnterCriticalSection (m_lock);

        T             * prev = (T *) item->m_prev;
        T             * next = (T *) item->m_next;
        if (prev == 0) {
                m_head = next;
        } else
                prev->m_next = next;
        if (next != 0)
                next->m_prev = prev;

        LeaveCriticalSection (m_lock);

        if (free)
                delete item;
}

/**
 * Remove and free an item based on a key.
 */

template <class T>
void SocketList<T> :: remove (SOCKET handle) {
        EnterCriticalSection (m_lock);

        T            ** link = & m_head;
        T             * scan = m_head;
        T             * prev = 0;
        while ((scan = * link) != 0) {
                if (scan->m_handle == handle) {
                        * link = (T *) scan->m_next;
                        scan->m_prev = prev;
                        break;
                }

                prev = scan;
                link = (T **) & scan->m_next;
        }

        LeaveCriticalSection (m_lock);

        if (scan != 0)
                delete scan;
}

/**
 * Structure for representing a replacement context.
 */

struct Replacement : public SocketTrack {
        unsigned long   m_length;
        unsigned long   m_offset;

        unsigned char * m_data;

                        Replacement (SOCKET handle) : SocketTrack (handle) { }
};

/**
 * Structure for representing an context where we're discarding output.
 */

struct Discarding : public SocketTrack {
        unsigned long   m_length;

                        Discarding (SOCKET handle) : SocketTrack (handle) { }
};

/**
 * The global list of bound event handles for sockets.
 */

SocketList<SocketTrack> l_events;

/**
 * The global list of active replacement items.
 */

SocketList<Replacement> l_replace;

/**
 * Global list of sockets where we are discarding sent data.
 */

SocketList<Discarding>  l_discard;

/**
 * Root registry key in which replacement items are located.
 */

HKEY                    l_rootKey;

/**
 * Set the root path used as the context for the names of replacement items.
 *
 * It would presumably be nice to support either or both of registry and file
 * paths here, so that users could choose. In the first instance I'll probably
 * only implement registry-hosted item support, but it would be sensible to
 * have the filesystem path setup present in the hosting machinery.
 */

void g_initReplacement (ReplaceHKEY key, const wchar_t * regPath) {
        LSTATUS         status = ERROR_NOT_FOUND;
        HKEY            result = 0;
        if (regPath != 0) {
                status = RegOpenKeyExW ((HKEY) key, regPath, 0, KEY_READ,
                                        & result);
        }

        if (status == ERROR_SUCCESS)
                l_rootKey = result;
}

/**
 * Do any unload-time cleanup.
 */

void g_unloadReplacement (void) {
        if (l_rootKey != 0) {
                RegCloseKey (l_rootKey);
                l_rootKey = 0;
        }

        l_replace.free ();
        l_events.free ();
}

/**
 * Add tracking for an event handle bound to a socket.
 *
 * This might potentially be used to change a binding from one event or even to
 * remove a binding, although most socket client applications don't do that.
 */

void g_addEventHandle (SOCKET handle, WSAEVENT event) {
        SocketTrack   * item = l_events.find (handle);
        if (item != 0) {
                item->m_event = event;
                return;
        }

        item = new SocketTrack (handle);
        l_events.add (* item);
}

/**
 * When a socket handle is being closed, remove any tracking data for it.
 */

void g_removeTracking (SOCKET handle) {
        l_events.remove (handle);
        l_replace.remove (handle);
        l_discard.remove (handle);
}

/**
 * Add a potential replacement document to the replacement set.
 *
 * This is a hook into us performed during rule parsing, although if we're
 * sourcing data from the registry we will probably do nothing.
 *
 * Whether file or registry is preferred, the name of the source is always
 * going to be relative to something; another problem with file access is that
 * the source directory for such content is unlikely to be one we can arrange
 * as such in the context of Steam itself, and it will need to have been set up
 * for us during the filter load somehow.
 */

void g_replacementCache (const wchar_t * /* name */) {
}

/**
 * Format the current local date and time as an RFC 822/RFC 1123 string.
 */

bool l_formatDate (char * buffer, size_t length) {
        if (length < 32)
                return false;

static  char          * days [7] = {
                "Sun", "Mon", "Tue", "Wed", "Thu", "Fri", "Sat"
        };
static  char          * months [13] = {
                "",
                "Jan", "Feb", "Mar", "Apr", "May", "Jun",
                "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"
        };

        SYSTEMTIME      now;
        GetSystemTime (& now);

        wsprintfA (buffer, "%s, %d %s %04d %02d:%02d:%02d GMT",
                   days [now.wDayOfWeek], now.wDay, months [now.wMonth],
                   now.wYear, now.wHour, now.wMinute, now.wSecond);

        return true;
}

/**
 * Helper for g_addReplacement (), create a replacement item record.
 */

bool g_addReplacement (SOCKET handle, const wchar_t * replacement, int status,
                       const char * extraText) {
        /*
         * Compute the space required to hold a UTF-8 version of the input, and
         * any template-type substitution needed (not that we support that as
         * yet).
         */

        int             utf8 = 0;
        if (replacement != 0) {
                utf8 = WideCharToMultiByte (CP_UTF8, 0, replacement, - 1, 0, 0, 0, 0);
                if (utf8 == 0)
                        return false;

                -- utf8;
        }

        /*
         * Given the size of the replacement document, we can prepare the HTTP
         * headers for it and thus know the size of the header and thus the
         * total combined space to allocate.
         */

        char            date [80];
        if (! l_formatDate (date, ARRAY_LENGTH (date)))
                return false;

        char            redirect [512];

        const char    * location = "";
        const char    * statusText = "UNKNOWN";
        switch (status) {
        case 200:
                if (replacement == 0 && extraText != 0)
                        utf8 = strlen (extraText);

                statusText = "OK";
                break;

        case 302:
                if (extraText == 0)
                        return false;

                wsprintfA (redirect, "Location: %s\r\n", extraText);
                location = redirect;

                statusText = "REDIRECT";
                break;

        default:
                if (extraText != 0)
                        statusText = extraText;
                break;
        }

        char            header [1024];
        wsprintfA (header,
                   "HTTP/1.1 %d %s\r\n"
                   "Date: %s\r\n"
                   "Expires: %s\r\n"
                   "Content-Type: text/html; charset=UTF8\r\n"
                   "Content-Length: %ld\r\n"
                   "Connection: Keep-Alive\r\n"
                   "%s"
                   "\r\n", status, statusText, date, date, utf8, location);

        size_t          headerLength = strlen (header);

        unsigned long   size = sizeof (Replacement) + headerLength + utf8 + 1;

        /*
         * Allocate the replacement control structure, and copy the replacement
         * document text into it (converting to UTF-8).
         */

        void          * mem = Replacement :: operator new (size);
        if (mem == 0)
                return false;

        Replacement   * item = new (mem) Replacement (handle);
        item->m_data = (unsigned char *) (item + 1);

        memcpy (item->m_data, header, headerLength);

        if (replacement != 0) {
                utf8 = WideCharToMultiByte (CP_UTF8, 0, replacement, - 1,
                                            (LPSTR) item->m_data + headerLength,
                                            utf8, 0, 0);
                if (utf8 == 0) {
                        delete item;
                        return false;
                }
                -- utf8;
        } else if (utf8 > 0) {
                /*
                 * An alternative way to supply content via a #200 rule. In
                 * this mode, we use ~ as an escape for '\n';
                 */

                unsigned char * dest = item->m_data + headerLength;
                unsigned long   count = utf8;
                while (count > 0) {
                        unsigned char   ch = * extraText;
                        ++ extraText;

                        if (ch == '~')
                                ch = '\n';

                        * dest = ch;

                        ++ dest;
                        -- count;
                }
        }

        item->m_length = headerLength + utf8;
        item->m_offset = 0;

        l_replace.add (* item);

        /*
         * Look for a bound event handle for the owner socket, and signal it as
         * we're making read data available.
         */

        SocketTrack   * track = l_events.find (item->handle ());
        if (track != 0)
                SetEvent (track->m_event);

        return true;
}

/**
 * Add a named replacement document item to be substituted on the indicated handle.
 */

bool g_addReplacement (SOCKET handle, const char * name, const char * /* url */) {
        wchar_t         tempName [80];
        int             result;
        result = MultiByteToWideChar (CP_UTF8, 0, name, - 1,
                                      tempName, ARRAY_LENGTH (tempName));
        if (result == 0)
                return false;

        /*
         * Determine whether the named item exists, and what its size in UTF-16
         * characters is.
         */

        LSTATUS         status;
        unsigned long   length = 0;
        status = RegQueryValueExW (l_rootKey, tempName, 0, 0, 0, & length);
        if (status != ERROR_SUCCESS) {
                OutputDebugStringA ("HTTP replacement not found\r\n");
                return false;
        }

        wchar_t       * replacement;
        replacement = (wchar_t *) HeapAlloc (GetProcessHeap (), 0, length);
        if (replacement == 0)
                return false;

        unsigned long   type;
        status = RegQueryValueExW (l_rootKey, tempName, 0, & type,
                                   (LPBYTE) replacement, & length);
        if (status != ERROR_SUCCESS) {
                HeapFree (GetProcessHeap (), 0, replacement);
                return false;
        }

        if (type == REG_MULTI_SZ) {
                /*
                 * For MULTI_SZ just join all the component strings by changing
                 * the interior terminator bytes into newlines.
                 */

                length = length / sizeof (wchar_t);
                wchar_t       * scan = replacement;
                wchar_t       * end = scan + length - 1;
                for (; scan != end ; ++ scan)
                        if (* scan == 0)
                                * scan = '\n';
        } else if (type == REG_SZ || type == REG_EXPAND_SZ) {
                length = length / sizeof (wchar_t);
        } else {
                HeapFree (GetProcessHeap (), 0, replacement);
                return false;
        }

        /*
         * Since a registry value has the potential to not be terminated,
         * ensure that a terminator is present.
         */

        if (replacement [length - 1] != 0) {
                replacement [length] = 0;
                ++ length;
        }

        bool            value;
        value = g_addReplacement (handle, replacement, 200, 0);
        HeapFree (GetProcessHeap (), 0, replacement);
        return value;
}

/**
 * Attempt to find a replacement item (and optionally, release it).
 */

Replacement * g_findReplacement (SOCKET handle) {
        return l_replace.find (handle);
}

/**
 * Given a replacement item, consume part of it for the caller.
 *
 * This is the simplest signature; it's up to the caller to adapt the incoming
 * API format to suit this (if we wanted to support multiple WSABUF structures,
 * for example, although for now we don't).
 *
 * Note that we're assuming here that a given source socket is only being read
 * on a single thread, so there is no need for locking here. That's probably
 * not a safe assumption in general to make if we ever get used with clients
 * written for high performance, but it's true of simple things like Steam.
 */

bool g_consumeReplacement (Replacement * item, unsigned long length, void * buf,
                           unsigned long * copied) {
        if (item == 0 || buf == 0)
                return false;

        unsigned long   avail = item->m_length - item->m_offset;
        if (length < avail)
                avail = length;

        memcpy (buf, item->m_data, avail);

        if (copied != 0)
                * copied = avail;

        item->m_offset += avail;

        if (item->m_offset != item->m_length)
                return true;

        /*
         * The replacement item has been consumed, remove it.
         */

        l_replace.remove (item, true);
        return true;
}

/**
 * Add a notice to discard data from a handle
 */

bool g_addDiscard (SOCKET handle, unsigned long length) {
        if (length == 0)
                return false;

        /*
         * Allocate the replacement control structure, and copy the replacement
         * document text into it (converting to UTF-8).
         */

        Discarding    * item = new Discarding (handle);
        if (item == 0)
                return false;

        item->m_length = length;
        l_discard.add (* item);
        return true;
}

/**
 * Determine if we're discarding output from this socket.
 */

Discarding * g_findDiscard (SOCKET handle) {
        return l_discard.find (handle);
}

/**
 * Consume some amount of data from the discard record.
 */

bool g_consumeDiscard (Discarding * item, unsigned long length, unsigned long * skip) {
        unsigned long   used = item->m_length;
        if (length < used)
                used = length;

        item->m_length -= used;
        length -= used;

        if (skip != 0)
                * skip = used;

        if (item->m_length == 0)
                l_discard.remove (item, true);
        return true;
}

/**@}*/
